from langchain.embeddings import CacheBackedEmbeddings
from langchain_community.embeddings import DashScopeEmbeddings
import time
from langchain.storage import LocalFileStore

#初始化模型
ali_embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",
    max_retries=3,
    dashscope_api_key="sk-4f1498f1c0314ba79ea2919bd7a02c4d"
)

storage = LocalFileStore("./embedding_cache/") #本地目录
cached_embeddings = CacheBackedEmbeddings.from_bytes_store(
    ali_embeddings,
    storage,
    namespace="openai_emb" #命名空间隔离不同模型
)

texts = ["小滴课堂AI大模型开发实战","小滴课堂AI大模型开发实战如何进行学习"] #故意重复

start_time = time.time()

#第一次使用，写入缓存
# emb1 = cached_embeddings.embed_documents(texts)
# print(f"首次调用嵌入维度:{len(emb1[0])}")
#
# embedded1_end_time = time.time()
# print(f"首次调用耗时:{embedded1_end_time - start_time}")
#
# #第二次调用（读取缓存）
# emb2 = cached_embeddings.embed_documents(texts)
# print(f"第二次调用结果相等:{emb1==emb2}")
#
# embedded2_end_time = time.time()
# print(f"二次调用耗时:{embedded2_end_time - embedded1_end_time}")


#场景二，使用embed_query不读取缓存

#第一次使用，写入缓存
emb1 = cached_embeddings.embed_query("小滴课堂AI大模型开发实战")
print(f"首次调用嵌入维度:{len(emb1)}")

embedded1_end_time = time.time()
print(f"首次调用耗时:{embedded1_end_time - start_time}")

#第二次调用（读取缓存）
emb2 = cached_embeddings.embed_query("小滴课堂AI大模型开发实战")
print(f"第二次调用结果相等:{emb1==emb2}")

embedded2_end_time = time.time()
print(f"二次调用耗时:{embedded2_end_time - embedded1_end_time}")